from cv2 import mean
import torch
import argparse
from torch._C import dtype
from torch.optim import optimizer
from torch.utils.data import DataLoader
from torch import nn, optim
from torchvision.transforms import transforms
from dataPro_origin import cellDate
import torch.nn.functional as F
from unet import Unet
import torch, gc
from save_re import Logger
import os 
from loss_p import dice_coef,iou_score
from tqdm import tqdm,trange

os.environ['CUDA_VISIBLE_DEVICES']='0'
trans = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize([0.5], [0.5])])
target_transform=transforms.ToTensor()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
def val_model(model,crterion,dataload,best_dice,best_iou):
    model=model.eval()
    print('val')
    with torch.no_grad():
        step=0
        epoch_loss = 0
        epoch_dice=0
        epoch_lou=0
        dt_size = len(dataload.dataset)
        opterizer = optim.Adam(model.parameters())
        for image,label in dataload:
            step+=1
            inputs=image.to(device)
            label=label.to(device)
            opterizer.zero_grad()
            output=model(inputs)
            loss=crterion(output,label)
            dice=dice_coef(output,label)
            lou=iou_score(output,label)
            # print("val,%d/%d,val_loss:%0.4f,dice_loss:%0.3f,Iou_score:%0.3f" % (step, (dt_size - 1) // dataload.batch_size + 1, loss.item(),dice.item(),lou.item())) 
            epoch_loss+=loss.item()
            epoch_dice+=dice.item()
            epoch_lou+=lou.item()
        print("val,loss:%0.4f dice:%0.3f lou:%0.3f" % (epoch_loss/step,epoch_dice/step,epoch_lou/step))
        mean_loss=epoch_loss/step
        mean_dice=epoch_dice/step
        mean_iou=epoch_lou/step
        if mean_dice>best_dice:
            print('mean_dice:{}>best_dice:{}'.format(mean_dice,best_dice))
            print('=======>save best dice model!')
            best_dice = mean_dice
            torch.save(model.state_dict(),r'dice_best'+'.pth')
        if mean_iou>best_iou:
            print('mean_iou:{}>best_iou:{}'.format(mean_iou,best_iou))
            print('=======>save best iou model!')
            best_iou=mean_iou
            torch.save(model.state_dict(),r'Iou_best'+'.pth')
        return best_dice,best_iou
def train_model(model,crterion,dataload,val_loader,num_epochs=301):
    best_loss = float('inf')
    best_dice=0
    best_iou=0
    for epoch in trange(num_epochs-1):
        model=model.train()
        print('train')
        print('Epoch {}/{}'.format(epoch, num_epochs - 1))
        print('-' * 10)
        dt_size = len(dataload.dataset)
        epoch_loss = 0
        epoch_dice=0
        epoch_iou=0
        step = 0
        opterizer = optim.Adam(model.parameters())
        for image,label in tqdm(dataload):
            step+=1
            inputs=image.to(device)
            label=label.to(device)
            opterizer.zero_grad()
            output=model(inputs)
            loss=crterion(output,label) 
            dice=dice_coef(output,label)
            iou=iou_score(output,label)
            # print("train,%d/%d,train_loss:%0.4f,dice_loss:%0.3f,Iou_score:%0.3f" % (step, (dt_size - 1) // dataload.batch_size + 1, loss.item(),dice.item(),iou.item())) 
            loss.backward() 
            opterizer.step()  
            epoch_loss+=loss.item()
            epoch_dice+=dice.item()
            epoch_iou+=iou.item() 
            gc.collect()
            torch.cuda.empty_cache()
        print("epoch %d loss:%0.4f dice:%0.3f lou:%0.3f" % (epoch, epoch_loss/step,epoch_dice/step,epoch_iou/step))
        # torch.save(model.state_dict(), 'Unet_weights_path58%d.pth' % epoch)
        best_dice,best_iou=val_model(model,crterion,val_loader,best_dice,best_iou)
    return model
if __name__ == "__main__":
    net = Unet(3, 1)
    # 将网络拷贝到deivce中
    net.to(device=device)
    # 指定训练址，开始训练
    train_path = r'cellseg_test'
    val_path=r'cellV'
    tran_set=cellDate(train_path, transform = trans,target_transform=target_transform)
    train_loader=torch.utils.data.DataLoader(dataset=tran_set,
                                               batch_size=2, 
                                               shuffle=True)
    val_set=cellDate(val_path, transform = trans,target_transform=target_transform)
    val_loader=torch.utils.data.DataLoader(dataset=val_set,
                                            batch_size=2,
                                            shuffle=True)
    criterion=nn.BCEWithLogitsLoss()
    train_model(net,criterion,train_loader,val_loader)      